from mindspore import Tensor
import numpy as np
import pytest
from mindspore import context # 记得加这个
from ..share.ops.primitive.rmsprop_ops import RmsPropMock

# 根据tensorflow.raw_ops文档描述，该算子不需要反向grad计算


'''
TEST_SUMMARY:test operator rmsprop, given (input_shape=(4, 2, 5, 3, 2, 4), dtype=np.float64)
'''
def test_rmsprop_input_6D_dtype_float64():
    var = np.ones([4, 2, 5, 3, 2, 4]).astype(np.float64)
    mean_square = np.ones([4, 2, 5, 3, 2, 4]).astype(np.float64)
    moment = np.ones([4, 2, 5, 3, 2, 4]).astype(np.float64)
    grad = np.ones([4, 2, 5, 3, 2, 4]).astype(np.float64)

    lr = Tensor(np.array([0.01]).astype(np.float64))
    decay = Tensor(np.array([0.0]).astype(np.float64))
    momentum = Tensor(np.array([1e-10]).astype(np.float64))
    epsilon = Tensor(np.array([0.001]).astype(np.float64))
    fact = RmsPropMock(var, mean_square, moment, grad, lr, decay, momentum, epsilon)
    if context.get_context('mode') == context.PYNATIVE_MODE:
        with pytest.raises(TypeError):
            fact.forward_cmp()
    else :
        fact.forward_cmp()


'''
TEST_SUMMARY:test operator rmsprop, given (input_shape=(2, 3, 3, 2, 4, 2, 4), dtype=np.complex64)
'''
def test_rmsprop_input_7D_dtype_complex64():
    var = np.ones([2, 3, 3, 2, 4, 2, 4]).astype(np.complex64)
    mean_square = np.ones([2, 3, 3, 2, 4, 2, 4]).astype(np.complex64)
    moment = np.ones([2, 3, 3, 2, 4, 2, 4]).astype(np.complex64)
    grad = np.ones([2, 3, 3, 2, 4, 2, 4]).astype(np.complex64)

    lr = Tensor(np.array([0.01 + 0.01j]).astype(np.complex64))
    decay = Tensor(np.array([0.0 + 0.0j]).astype(np.complex64))
    momentum = Tensor(np.array([1e-10 + 1e-10j]).astype(np.complex64))
    epsilon = Tensor(np.array([0.001 + 0.001j]).astype(np.complex64))
    fact = RmsPropMock(var, mean_square, moment, grad, lr, decay, momentum, epsilon)
    if context.get_context('mode') == context.PYNATIVE_MODE:
        with pytest.raises(TypeError):
            fact.forward_cmp()
    elif context.get_context('device_target') == 'CPU':
        with pytest.raises(TypeError):
            fact.forward_cmp()
    else:
        fact.forward_cmp()